import torch
import numpy as np

from megatron import print_rank_0

def dummy_datasets_provider(data_prefix, data_impl, splits_string,
                                    train_valid_test_num_samples,
                                    seq_length, seed, skip_warmup):
    num_samples=train_valid_test_num_samples[0]
    dataset=DummyDataset(num_items=num_samples, item_size=seq_length)
    return dataset, None, None

class DummyDataset(torch.utils.data.Dataset):
    def __init__(self, num_items, item_size):
        super().__init__()
        self.num_items = num_items
        self.item_size = item_size

    def __getitem__(self, index):
        return {'text': np.arange(self.item_size+1, dtype=np.int64)}

    def __len__(self):
        return self.num_items

    @property
    def sizes(self):
        return np.array([self.item_size] * self.num_items)

    def num_tokens(self, index):
        return self.item_size

    def size(self, index):
        return self.item_size

    @property
    def supports_prefetch(self):
        return False
